#!/usr/bin/env python3
"""
重建ChromaDB数据，包含metadata信息
"""

import sys
import os
import django

# 设置Django环境
sys.path.append('/Users/baimu/PycharmProjects/2505A/boss-llm/boss')
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'boss.settings')
django.setup()

def rebuild_chroma_with_metadata():
    """
    重建ChromaDB数据，包含完整的metadata信息
    """
    print("=" * 80)
    print("重建ChromaDB数据，包含metadata")
    print("=" * 80)
    
    try:
        from rag.job_recommendation.job_recommendation import fetch_job, build_searchable_text_from_job, job_with_text_embedding, sync_jobs_to_chroma
        import chromadb
        
        print("1. 获取岗位数据...")
        job_list = fetch_job(limit=200)  # 获取更多数据
        print(f"   获取到 {len(job_list)} 个岗位")
        
        print("2. 构建搜索文本...")
        job_text_list = []
        for job in job_list:
            job_str = build_searchable_text_from_job(job)
            job_text_list.append({
                'id': job['id'],
                'title': job['title'],
                'company': job['company'],
                'city': job['city'],
                'tags': job.get('tags', ''),
                'education': job.get('education', ''),
                'working_years': job.get('working_years', ''),
                'description': job.get('description', ''),
                'label': job.get('label', ''),
                'job_str': job_str
            })
        
        print("3. 生成文本嵌入...")
        job_text_embedding_list = job_with_text_embedding(job_text_list)
        print(f"   生成 {len(job_text_embedding_list)} 个嵌入向量")
        
        print("4. 删除旧集合...")
        try:
            client = chromadb.PersistentClient(path="./chromadb_db")
            client.delete_collection(name="job_text_embedding_0927")
            print("   旧集合已删除")
        except Exception as e:
            print(f"   删除旧集合时出错: {str(e)}")
        
        print("5. 存储到ChromaDB（包含metadata）...")
        sync_jobs_to_chroma(job_text_embedding_list)
        print("   数据存储完成")
        
        print("6. 验证存储结果...")
        collection = client.get_collection(name="job_text_embedding_0927")
        count = collection.count()
        print(f"   集合中文档数量: {count}")
        
        # 检查metadata
        if count > 0:
            sample = collection.get(limit=1, include=["metadatas"])
            if sample and sample.get("metadatas") and sample["metadatas"][0]:
                metadata = sample["metadatas"][0]
                print(f"   样本文档metadata: {metadata}")
                print("   ✅ metadata存储成功")
            else:
                print("   ❌ metadata存储失败")
        
        print("\n🎉 ChromaDB数据重建完成！")
        
    except Exception as e:
        print(f"❌ 重建失败: {str(e)}")
        import traceback
        traceback.print_exc()


def test_rebuilt_data():
    """
    测试重建后的数据
    """
    print("\n" + "=" * 80)
    print("测试重建后的数据")
    print("=" * 80)
    
    try:
        from rag.job_recommendation.job_recommendation import search_job_by_chroma
        
        # 测试查询
        query_text = "Java开发工程师"
        
        print(f"测试查询: {query_text}")
        results = search_job_by_chroma(query_text, top_k=3)
        
        if results:
            print(f"检索到 {len(results)} 个结果:")
            for i, job in enumerate(results, 1):
                print(f"{i}. {job.get('title', 'N/A')} - {job.get('company', 'N/A')} - {job.get('city', 'N/A')}")
            
            # 检查metadata
            has_metadata = any(job.get('title', '').strip() for job in results)
            print(f"\nMetadata状态: {'✅ 有数据' if has_metadata else '❌ 无数据'}")
        else:
            print("❌ 没有检索到结果")
            
    except Exception as e:
        print(f"❌ 测试失败: {str(e)}")


def main():
    """
    主函数
    """
    print("开始重建ChromaDB数据...")
    
    rebuild_chroma_with_metadata()
    test_rebuilt_data()
    
    print("\n重建完成！")


if __name__ == "__main__":
    main()
